from langchain.prompts import MessagesPlaceholder
from langchain_deepseek import ChatDeepSeek
from langchain.memory import ConversationBufferMemory
from langchain.agents import initialize_agent, Tool
from langchain.agents import AgentType
from langchain.agents.tools import tool
from langchain.globals import set_debug
import os

@tool
def add(a:int) -> int:
    """Add two numbers together"""
    return a + 1

class SearchAPI:
    def __init__(self) -> None:
        try:
            api_key = os.environ['API_KEY']
            api_base = os.environ['API_BASE']
            model_name = os.environ['MODEL_NAME']
        except KeyError:
            raise ValueError('DEEPSEEK_API_KEY未在环境变量中配置')

        set_debug(True)
        
        llm = ChatDeepSeek(
            model=model_name,
            api_key=api_key,
            api_base=api_base,
            timeout=1800,
        )
        agent_kwargs = {
            "extra_prompt_messages": [MessagesPlaceholder(variable_name="memory")],
        }
        memory = ConversationBufferMemory(memory_key="memory", return_messages=True)

        self.agent = initialize_agent(
            llm=llm, 
            agent=AgentType.OPENAI_FUNCTIONS, 
            tools=[add], 
            verbose=True, 
            agent_kwargs=agent_kwargs, 
            memory=memory
        )

    def get_serach(self,msg: str) -> str:

        return self.agent.run(msg)